import rclpy
from rclpy.node import Node
from std_msgs.msg import String
from smbus2 import SMBus
import json

class PCA9685:
    def __init__(self, address=0x40, bus_number=1):
        self.address = address
        self.bus = SMBus(bus_number)
        self.setup_pwm()

    def setup_pwm(self):
        # 初始化PWM为50Hz（舵机标准频率）
        self.bus.write_byte_data(self.address, 0x00, 0x10)  # 睡眠模式
        prescale = int(25000000.0 / 4096.0 / 50.0 - 1)
        self.bus.write_byte_data(self.address, 0xFE, prescale)
        self.bus.write_byte_data(self.address, 0x00, 0x80)  # 唤醒

    def set_servo_pulse(self, channel, pulse):
        # 设置舵机脉宽（0-4095）
        self.bus.write_byte_data(self.address, 0x06 + 4 * channel, 0 & 0xFF)
        self.bus.write_byte_data(self.address, 0x07 + 4 * channel, 0 >> 8)
        self.bus.write_byte_data(self.address, 0x08 + 4 * channel, pulse & 0xFF)
        self.bus.write_byte_data(self.address, 0x09 + 4 * channel, pulse >> 8)

class ServoControlNode(Node):
    def __init__(self):
        super().__init__('servo_control_node')
        # 订阅统一话题，接收所有指令（WebSocket和串口）
        self.subscription = self.create_subscription(
            String,
            'websocket_cmd',
            self.handle_unified_cmd,
            10
        )
        
        # 初始化两个PCA9685模块（地址0x40和0x41）
        try:
            self.pca_0x40 = PCA9685(0x40, 1)
            self.pca_0x41 = PCA9685(0x41, 1)
            self.get_logger().info("舵机控制节点启动，已订阅: websocket_cmd")
        except Exception as e:
            self.get_logger().error(f"PCA9685初始化失败: {e}")
            self.pca_0x40 = None
            self.pca_0x41 = None

    def handle_unified_cmd(self, msg):
        """处理统一格式的JSON指令"""
        self.get_logger().info(f"收到指令: {msg.data}")
        try:
            # 解析JSON格式的指令列表
            cmd_list = json.loads(msg.data)
            if not isinstance(cmd_list, list):
                self.get_logger().error("指令格式错误，应为列表类型")
                return
            
            # 处理每条指令
            for cmd in cmd_list:
                if not isinstance(cmd, dict):
                    self.get_logger().error(f"无效指令格式: {cmd}（应为字典）")
                    continue
                
                # 提取必要字段
                board_id = cmd.get("b")
                channel = cmd.get("c")
                pulse = cmd.get("p")
                
                # 检查字段完整性
                if not all([board_id is not None, channel is not None, pulse is not None]):
                    self.get_logger().error(f"指令缺少字段: {cmd}（需包含b, c, p）")
                    continue
                
                # 转换为整数
                try:
                    board_id = int(board_id)
                    channel = int(channel)
                    pulse = int(pulse)
                except ValueError:
                    self.get_logger().error(f"指令数值错误: {cmd}（b, c, p应为数字）")
                    continue
                
                # 执行舵机控制
                self.control_servo(board_id, channel, pulse)
        
        except json.JSONDecodeError:
            self.get_logger().error(f"指令解析失败: {msg.data}（不是有效的JSON）")
        except Exception as e:
            self.get_logger().error(f"处理指令时出错: {str(e)}")

    def control_servo(self, board_id, channel, pulse):
        """根据板ID和通道控制舵机"""
        # 选择对应的PCA9685模块
        if board_id == 0:
            pca = self.pca_0x40
        elif board_id == 1:
            pca = self.pca_0x41
        else:
            self.get_logger().error(f"不支持的板ID: {board_id}（仅支持0或1）")
            return
        
        # 检查模块是否初始化成功
        if pca is None:
            self.get_logger().error(f"板{board_id}未初始化，无法控制")
            return
        
        # 检查通道范围（0-15）
        if not (0 <= channel <= 15):
            self.get_logger().error(f"通道超出范围: {channel}（需0-15）")
            return
        
        # 限制脉宽范围（100 - 599 对应舵机0~180度）
        pulse = max(100, min(599, pulse))
        
        # 发送控制指令
        try:
            pca.set_servo_pulse(channel, pulse)
            self.get_logger().info(f"控制成功: 板{board_id} 通道{channel} 脉宽{pulse}")
        except Exception as e:
            self.get_logger().error(f"控制舵机失败: {str(e)}")

def main(args=None):
    rclpy.init(args=args)
    node = ServoControlNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info("用户中断，退出节点")
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()